"""
python fujian2_RandomForest_fix.py
"""

import os
import json
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

def RandomForest_fillBlank_single_category(input_json_path, output_json_path):
    # 创建输出目录
    output_folder = os.path.dirname(output_json_path)
    graph_folder = os.path.join(output_folder, 'graph')
    os.makedirs(output_folder, exist_ok=True)
    os.makedirs(graph_folder, exist_ok=True)
    
    # 读取输入数据
    with open(input_json_path, 'r') as f:
        data = json.load(f)
    
    # 转换数据为DataFrame
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    
    # 生成完整的日期范围
    start_date, end_date = "2022-07-01", "2023-06-30"
    date_range = pd.date_range(start=start_date, end=end_date)
    
    # 插入空白日期并标记缺失值
    full_df = pd.DataFrame(index=date_range)
    full_df = full_df.join(df['sales'])
    full_df['is_missing'] = full_df['sales'].isna()
    
    # 提取已有数据训练随机森林
    train_df = full_df[~full_df['is_missing']]
    X_train = (train_df.index - train_df.index[0]).days.values.reshape(-1, 1)
    y_train = train_df['sales'].values
    
    # 使用随机森林回归进行预测
    model = RandomForestRegressor(n_estimators=100, random_state=0)
    model.fit(X_train, y_train)
    
    # 填补中间空白段的时间
    missing_df = full_df[full_df['is_missing']]
    X_missing = (missing_df.index - full_df.index[0]).days.values.reshape(-1, 1)
    full_df.loc[full_df['is_missing'], 'sales'] = model.predict(X_missing)
    
    # 将填补后的数据保存至JSON文件
    filled_data = full_df.loc["2022-10-01":"2023-03-31"]
    filled_data = filled_data[['sales']].reset_index()
    filled_data.columns = ['date', 'sales']
    filled_data['date'] = filled_data['date'].dt.strftime('%Y/%m/%d')
    
    with open(output_json_path, 'w') as f:
        json.dump(filled_data.to_dict(orient='records'), f, indent=4)
    
    # 绘制图表，显示原始和填补后的数据
    plt.figure(figsize=(12, 6))
    plt.plot(full_df.index, full_df['sales'], label='填补后的数据', color='orange')
    plt.plot(train_df.index, train_df['sales'], label='原始数据', color='blue')
    plt.xlabel('日期')
    plt.ylabel('销售量')
    plt.legend()
    plt.title('销售数据填补')
    
    # 保存图表
    graph_path = os.path.join(graph_folder, 'sales_data_fill.png')
    plt.savefig(graph_path)
    plt.close()
    
    print(f"填补后的数据已保存至 {output_json_path}")
    print(f"图表已保存至 {graph_path}")

if __name__ == '__main__':
    input_path = "../fujian/fujian2/groupByCategory/category_category1.json"
    output_path = "../fujian/fujian2/RandomForest_fix/category_category1.json"
    RandomForest_fillBlank_single_category(input_path, output_path)